package com.bcw.springtest.config;

import org.springframework.beans.BeansException;
import org.springframework.beans.factory.BeanFactory;
import org.springframework.beans.factory.BeanFactoryAware;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.web.accept.ContentNegotiationManager;
import org.springframework.web.method.support.HandlerMethodReturnValueHandler;
import org.springframework.web.servlet.mvc.method.annotation.RequestMappingHandlerAdapter;
import org.springframework.web.servlet.mvc.method.annotation.RequestResponseBodyMethodProcessor;

import java.util.ArrayList;
import java.util.List;

public class BcwRequestResponseMappingHandlerAdapter implements BeanFactoryAware, InitializingBean {

    private BeanFactory beanFactory;

    @Override
    public void setBeanFactory(BeanFactory beanFactory) throws BeansException {
        this.beanFactory = beanFactory;
    }

    @Override
    public void afterPropertiesSet() {

        RequestMappingHandlerAdapter bean = beanFactory.getBean(RequestMappingHandlerAdapter.class);

        List<HandlerMethodReturnValueHandler> newList = new ArrayList<>();

        for (HandlerMethodReturnValueHandler handler : bean.getReturnValueHandlers()) {
            if (handler instanceof RequestResponseBodyMethodProcessor) {
                newList.add(new BcwResponseBodyWrapper(bean.getMessageConverters(),
                        beanFactory.getBean("mvcContentNegotiationManager", ContentNegotiationManager.class),
                        beanFactory));
            } else {
                newList.add(handler);
            }
        }
        bean.setReturnValueHandlers(newList);
    }
}
